import baostock as bs
import pandas as pd
from datetime import datetime,timedelta
import os
from utils import read_config,get_directory
from logger import logger


# 读取配置信息
config = read_config()
# A交易日（放量）前交易天数
preday_num = config.getint('volume_increase', 'preday_num');
# A交易日（放量）倍数
quantity_multiplier = config.getfloat('volume_increase', 'quantity_multiplier')
# A交易日（放量） 前后平均放量倍数
average_multiple = config.getfloat('volume_increase', 'average_multiple')
# 最后一个交易日是前期地点涨幅不超过倍数
last_day_price_multiple = config.getfloat('volume_increase', 'last_day_price_multiple')
# 考察周期，默认最后30个交易日
last_day_num = config.getint('volume_increase', 'last_day_num')

start_day = config.get('settings','daily_start_day')
_base_directory = config.get('settings','base_file')
_daily_directory = config.get('settings','daily_file')
_volume_file = config.get('settings','volume_file')

daily_directory = get_directory(_base_directory, _daily_directory)
# 使用正则表达式提取目标子字符串
_file_name_pattern = r'(sz|sh)\.\d+'


def formate_code(stock_code):
    if stock_code.startswith('6') and len(stock_code) == 6:
        return 'sh.'+stock_code
    elif stock_code.startswith('0') and len(stock_code) == 6:
        return 'sz.'+stock_code
    elif stock_code.startswith('3') and len(stock_code) == 6:
        return 'sz.'+stock_code
    else:
        return stock_code


def check_daily_all():
    rs = []
    # 获取目录中的所有文件
    files = os.listdir(daily_directory)
    # 如果目录为空，直接返回
    if not files:
        logger.warn(f"目录 {daily_directory} 下没有文件.")
        return rs
    # 遍历目录中的所有文件
    for filename in files:
        # 拼接文件的完整路径
        file_path = os.path.join(daily_directory, filename)
        _row = _check_daily_file(file_path)
        if _row is not None:
            rs.append(_row)
    logger.info(rs)
    return rs


def get_daily_path_name(code):
    return os.path.join(daily_directory,'daily_'+formate_code(code)+'.csv')


def check_daily(code):
    file_name = get_daily_path_name(code)
    return _check_daily_file(file_name)


def _check_daily_file(file_name):
    # 从CSV文件加载数据
    df = pd.read_csv(file_name, parse_dates=['date'])
    # 按日期排序
    df.sort_values('date', inplace=True)
    # 计算每日成交量均值
    df['avg_volume'] = df['volume'].rolling(window=30).mean()
    # 计算振幅
    df['amplitude'] = (df['high'] - df['low']) / df['open'] * 100
    # df.to_csv('./data/daily_' + code + '-bak.csv', index=False, encoding='utf-8-sig')
    # 找出最后30个交易日
    last_30_days = df[-last_day_num:]
    last_close_price = df.iloc[-1]['close']
    # 遍历最后30个交易日
    for index, row in last_30_days.iterrows():
        # 找到A日前的交易日
        # previous_days = df.loc[:index]
        previous_days = df.loc[index - preday_num - 1:index - 1]
        # 检查A日成交量是否满足条件
        if row['volume'] > previous_days['volume'].max() * quantity_multiplier:
            # 检查A日后的交易日
            after_days = df.loc[index+2:]
            # 检查A日后平均成交量是否满足条件
            if after_days['volume'].mean() > row['avg_volume'] * average_multiple:
                if after_days['volume'].max() < row['volume']:
                    # 检查收盘价是否低于前50个交易日最低价的1.2倍
                    min_price_50 = previous_days['low'].min()
                    max_price_50 = previous_days['high'].max()
                    if last_close_price < min_price_50 * last_day_price_multiple and row['close'] < max_price_50*1.1:
                        # 筛选振幅小于6%的数据
                        filtered_df = after_days[after_days['amplitude'] < 6]
                        if len(filtered_df) / len(after_days) > 0.7:
                            # 检查A日是阳线
                            if row['close'] > row['open']:
                                # 去掉ST股
                                if row['isST'] == 0:
                                    # 满足条件，存储数据到文件
                                    logger.info(f"满足条件的股票：{row['code']}，第一个爆发日:{row['date']}")
                                    return [row['code'],row['date']]
    return None


# adjustflag：复权类型，默认不复权：3；1：后复权；2：前复权。已支持分钟线、日线、周线、月线前后复权。
def query_history_k(code,start_day = None,end_day =None,frequency="d",adjustflag="2" ):
    if start_day is None:
        start_day = (datetime.now() - timedelta(days=90)).strftime('%Y-%m-%d')
    if end_day is None:
        end_day = datetime.now().strftime('%Y-%m-%d')
    # 获取沪深A股历史K线数据
    rs = bs.query_history_k_data_plus(formate_code(code),
                                    "date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,isST",
                                    start_date =start_day,
                                    end_date =end_day,
                                    frequency =frequency,
                                    adjustflag =adjustflag)

    # 打印结果集
    data_list = []
    if(rs.error_code != '0'):
        logger.warn(rs.error_code)
    while (rs.error_code == '0') & rs.next():
        # 获取一条记录，将记录合并在一起
        data_list.append(rs.get_row_data())
    return pd.DataFrame(data_list, columns=rs.fields)


def get_previous_workday(date_str):
    # 将字符串转换为pandas的Timestamp对象
    date_obj = pd.Timestamp(date_str)
    # 如果当前日期是工作日（周一至周五），直接返回
    if date_obj.weekday() < 5:
        return date_str
    # 如果不是工作日，找到之前的最近工作日
    # 从当前日期开始，向前遍历，直到找到工作日
    while date_obj.weekday() >= 5:
        date_obj -= pd.Timedelta(days=1)  # 减去一天
    return date_obj.strftime('%Y-%m-%d')


def reload_daily_k(current_date = None, file_dir = _daily_directory, start_day = start_day):
    if current_date is None:
        # 获取当前日期
        # 格式化日期为 'YYYY-MM-DD' 格式
        current_date = datetime.now().strftime('%Y-%m-%d')
    # 登录系统
    lg = bs.login()
    logger.info('login respond error_code:'+lg.error_code)
    logger.info('login respond error_msg:'+lg.error_msg)
    daily_directory = get_directory(_base_directory, file_dir)
    rs = bs.query_all_stock(get_previous_workday(current_date)).get_data()
    for index, row in rs.iterrows():
        logger.info(f"Index: {index}")
        logger.info(f"code: {row['code']}")
        stock_code = row['code']
        if not (stock_code.startswith('bj.') or stock_code.startswith('sh.0') or stock_code.startswith('sz.39')):
        # if stock_code =='sh.600900' or stock_code == 'sz.301469':
            _file_path = os.path.join(daily_directory , 'daily_' + stock_code + '.csv')
            if os.path.exists(_file_path):
                # 从CSV文件加载数据
                _old_df = pd.read_csv(_file_path)
            else:
                _old_df = None
            if _old_df is not None and len(_old_df) > 0:
                # 按日期排序
                _old_df.sort_values('date', inplace=True)
                _date_obj = _old_df['date'].iloc[-1]
                _next_day = datetime.strptime(_date_obj, "%Y-%m-%d") + timedelta(days=1)
                _start_day = _next_day.strftime("%Y-%m-%d")
                _new_df = query_history_k(stock_code, start_day=_start_day)
                # 合并两个 DataFrame
                df = pd.concat([_old_df, _new_df], ignore_index=True)
                # 去除基于 'date' 列的重复行
                df = df.drop_duplicates(subset='date', keep='last')
            else:
                df = query_history_k(stock_code, start_day=start_day)
            df.to_csv(_file_path, index=False, encoding='utf-8-sig')
    # 登出系统
    bs.logout()
    logger.info("结束")


def get_daily_file(stock_code,selected_fields = None,tail_num = 30,file_dir = _daily_directory):
    daily_directory = get_directory(_base_directory, file_dir)
    file_name = daily_directory + '/daily_' + formate_code(stock_code) + '.csv'
    return get_daily_by_file(file_name, selected_fields, tail_num)


def get_daily_by_file(file_name, selected_fields = None, tail_num = None):
    if os.path.exists(file_name):
        # 从CSV文件加载数据
        df = pd.read_csv(file_name, parse_dates=['date'])
        # 按日期排序
        df.sort_values('date', inplace=True)
        if tail_num is not None:
            rs = df.tail(tail_num)
        else:
            rs = df
        if selected_fields is not None:
            # 提取所选字段
            return rs[selected_fields]
        return rs
    else:
        return None

    # 定义文件读取函数
def read_files(directory):
    # 获取目录中的所有文件
    files = os.listdir(directory)
    # 如果目录为空，直接返回
    if not files:
        logger.warn(f"目录 {daily_directory} 下没有文件.")
        return []
    all_files = []
    for filename in files:
        file_path = os.path.join(directory, filename)
        df = get_daily_by_file(file_path)
        all_files.append(df)
    return all_files

if __name__ == "__main__":
    check_daily('603960')
    # # 获取当前日期
    # current_date = datetime.now().date()
    # # 格式化日期为 'YYYY-MM-DD' 格式
    # formatted_date = current_date.strftime('%Y-%m-%d')
    #
    # # 登录系统
    # lg = bs.login()
    # print('login respond error_code:', lg.error_code)
    # print('login respond error_msg:', lg.error_msg)
    # rs = bs.query_all_stock(get_previous_workday(formatted_date)).get_data()
    # for index, row in rs.iterrows():
    #     print(f"Index: {index}")
    #     print(f"code: {row['code']}")
    #     stock_code = row['code']
    #     if not stock_code.startswith('bj.'):
    #         df = query_history_k(stock_code,start_day = '2023-01-01')
    #         df.to_csv('./data/daily_'+stock_code+'.csv', index=False, encoding='utf-8-sig')
    # # 登出系统
    # bs.logout()
    # print("结束")